#include "processUtil.h"
#include <fstream>
#include <algorithm>
#include <functional>
#include <sstream>

namespace Util
{

    FILE *ProcessUtil::PopenEX(const std::string &cmd, PID_T *pidOut)
    {
#if defined(__unix__) || defined(__APPLE__)
        /*
        使用fork+exec函数的组合，而不是直接popen的原因在于：
        popen也会创建子进程，子进程在tshark执行后会不断输出，从而导致子线程中的fgets一直有返回数据，
        无法结束（似乎pclose后也会阻塞），所以需要手动kill掉进程。而kill掉进程需要得到进程的pid，pid需要通过fork获得
        */
        int pipefd[2] = {0};
        FILE *pipeFd = nullptr;
        if (pipe(pipefd) == -1)
        {
            perror("pipe");
            return pipeFd;
        }

        pid_t pid = fork();
        if (pid == -1)
        {
            perror("fork");
            close(pipefd[0]);
            close(pipefd[1]);
            return pipeFd;
        }
        /*
        note: 多进程编程——父进程捕获子进程输出 importance!!!
            初始化管道：
                int pipefd[2] = {0};
                if(pipe(pipefd)==-1){
                    // error process...
                }
            创建进程 fork()：
                pid_t pid = fork();
                if(pid == -1){
                    // close pipefd
                    // error process...
                }
            子进程(pid == 0)：
                1、将 STDERR_FILENO、STDOUT_FILENO 重定向到写端，供父进程读取
                2、关闭读端、关闭写端，防止子进程直接使用原始管道
                3、调用exec相关函数，替换进程的内存映像
            父进程(pid == 子进程的pid)：
                1、处理读端子进程的数据，看情况是否需要使用写端与子进程进行通信
                2、处理数据
                3、等待子进程的回收 wait、waitpid相关函数
        */
        if (pid == 0) // child process
        {
            close(pipefd[0]);               // 关闭读端
            dup2(pipefd[1], STDERR_FILENO); // 将 stderr 重定向到管道
            dup2(pipefd[1], STDOUT_FILENO); // 将 stdout 重定向到管道
            close(pipefd[1]);               // 关闭管道写端，防止子进程使用原始的管道写端，而只能将stdout、stderr输出到写端中

            execl("/bin/sh", "sh", "-c", cmd.c_str(), NULL); // 执行命令
            _exit(1);                                        // execl执行失败
        }

        // 父进程将读管道，关闭写端
        close(pipefd[1]);
        pipeFd = fdopen(pipefd[0], "r");

        // 设置进程id
        if (pidOut)
        {
            *pidOut = pid;
        }
        return pipeFd;

#elif (_WIN32)
        HANDLE hReadPipe, hWritePipe;
        SECURITY_ATTRIBUTES saAttr;
        PROCESS_INFORMATION piProcInfo;
        STARTUPINFO siStartInfo;
        FILE *pipeFp = nullptr;

        // 设置安全属性，允许管道句柄继承
        saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
        saAttr.bInheritHandle = TRUE;
        saAttr.lpSecurityDescriptor = nullptr;

        // 创建匿名管道
        if (!CreatePipe(&hReadPipe, &hWritePipe, &saAttr, 0))
        {
            perror("CreatePipe");
            return nullptr;
        }

        // 确保写句柄不被子进程继承
        if (!SetHandleInformation(hReadPipe, HANDLE_FLAG_INHERIT, 0))
        {
            perror("SetHandleInformation");
            CloseHandle(hReadPipe);
            CloseHandle(hWritePipe);
            return nullptr;
        }

        // 初始化 STARTUPINFO 结构体
        ZeroMemory(&piProcInfo, sizeof(PROCESS_INFORMATION));
        ZeroMemory(&siStartInfo, sizeof(STARTUPINFO));
        siStartInfo.cb = sizeof(STARTUPINFO);
        siStartInfo.hStdError = hWritePipe;
        siStartInfo.hStdOutput = hWritePipe;
        siStartInfo.dwFlags |= STARTF_USESTDHANDLES;

        // 创建子进程
        if (!CreateProcess(
                nullptr,               // No module name (use command line)
                (LPSTR)command.data(), // Command line
                nullptr,               // Process handle not inheritable
                nullptr,               // Thread handle not inheritable
                TRUE,                  // Set handle inheritance
                CREATE_NO_WINDOW,      // No window
                nullptr,               // Use parent's environment block
                nullptr,               // Use parent's starting directory
                &siStartInfo,          // Pointer to STARTUPINFO structure
                &piProcInfo            // Pointer to PROCESS_INFORMATION structure
                ))
        {
            perror("CreateProcess");
            CloseHandle(hReadPipe);
            CloseHandle(hWritePipe);
            return nullptr;
        }

        // 关闭写端句柄（父进程不使用）
        CloseHandle(hWritePipe);

        // 返回子进程 PID
        if (pidOut)
        {
            *pidOut = piProcInfo.dwProcessId;
        }

        // 将管道的读端转换为 FILE* 并返回
        pipeFp = _fdopen(_open_osfhandle(reinterpret_cast<intptr_t>(hReadPipe), _O_RDONLY), "r");
        if (!pipeFp)
        {
            CloseHandle(hReadPipe);
        }

        // 关闭进程句柄（不需要等待子进程）
        CloseHandle(piProcInfo.hProcess);
        CloseHandle(piProcInfo.hThread);

        return pipeFp;

#endif
    }

    int ProcessUtil::Kill(PID_T pid)
    {
#if defined(__unix__) || defined(__APPLE__)
        return kill(pid, SIGTERM);
#elif (_WIN32)
        // 打开指定进程
        HANDLE hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        if (hProcess == nullptr)
        {
            std::cout << "Failed to open process with PID " << pid << ", error: " << GetLastError() << std::endl;
            return -1;
        }

        // 终止进程
        if (!TerminateProcess(hProcess, 0))
        {
            std::cout << "Failed to terminate process with PID " << pid << ", error: " << GetLastError() << std::endl;
            CloseHandle(hProcess);
            return -1;
        }

        // 成功终止进程
        CloseHandle(hProcess);
        return 0;
#endif
    }

    bool ProcessUtil::Exec(const std::string &cmd)
    {
#ifdef _WIN32
        PROCESS_INFORMATION piProcInfo;
        STARTUPINFO siStartInfo;

        // 初始化 STARTUPINFO 结构体
        ZeroMemory(&piProcInfo, sizeof(PROCESS_INFORMATION));
        ZeroMemory(&siStartInfo, sizeof(STARTUPINFO));

        // 创建子进程
        if (CreateProcess(
                nullptr,               // No module name (use command line)
                (LPSTR)cmdline.data(), // Command line
                nullptr,               // Process handle not inheritable
                nullptr,               // Thread handle not inheritable
                TRUE,                  // Set handle inheritance
                CREATE_NO_WINDOW,      // No window
                nullptr,               // Use parent's environment block
                nullptr,               // Use parent's starting directory
                &siStartInfo,          // Pointer to STARTUPINFO structure
                &piProcInfo            // Pointer to PROCESS_INFORMATION structure
                ))
        {
            WaitForSingleObject(piProcInfo.hProcess, INFINITE);
            CloseHandle(piProcInfo.hProcess);
            CloseHandle(piProcInfo.hThread);
            return true;
        }
        else
        {
            return false;
        }
#else
        return std::system(cmd.c_str()) == 0;
#endif
    }

    void ProcessUtil::getProcessInfo(pid_t pid, pid_t &parentPid, std::string &name)
    {
        std::string statusFilePath = "/proc/" + std::to_string(pid) + "/status";
        std::ifstream statusFile(statusFilePath);
        if (!statusFile.is_open())
        {
            LOG_F(ERROR, "failed in open file: %s\n", statusFilePath.c_str());
            return;
        }
        std::string line;
        while (std::getline(statusFile, line))
        {
            std::stringstream iss(line);
            std::string key, val;
            if (!(iss >> key >> val))
            {
                continue;
            }
            if (key == "Name:")
            {
                name = val;
            }
            if (key == "PPid:")
            {
                parentPid = std::stol(val);
            }
        }
    }

    void ProcessUtil::getProcessCmdline(pid_t pid, std::string &cmdline)
    {
        std::string cmdPath = "/proc/" + std::to_string(pid) + "/cmdline";
        std::ifstream cmdFile(cmdPath);
        if (!cmdFile.is_open())
        {
            LOG_F(ERROR, "failed in open file: %s\n", cmdPath.c_str());
            return;
        }
        /**
         * note: 将整个文件一次性读入的方法：
         *  1、rdbuf()
         *  2、按行读取（简单，如果不是对每行进行处理，性能不如以上两种方法）
         *  3、按块4096B读取
         * */

        // // 1、直接rdbuf读入整个文件
        // std::stringstream buffer;
        // buffer << cmdFile.rdbuf();
        // cmdline = buffer.str();

        // // 2、按行读取
        // std::string line;
        // while(std::getline(cmdFile,line)){
        //     cmdline.append(line);
        // }

        // 3、按块读取
        cmdline = "";
        std::vector<char> buffer(4096);
        while (cmdFile.read(buffer.data(), buffer.size()))
        {
            cmdline.append(buffer.data(), cmdFile.gcount());
        }
        cmdline.append(buffer.data(), cmdFile.gcount());

        // 将 \0 替换为 ' '
        std::replace(cmdline.begin(), cmdline.end(), '\0', ' ');
        if (!cmdline.empty())
        {
            // 去掉最后一个空格
            cmdline.pop_back();
        }
    }

    std::vector<pid_t> ProcessUtil::getChildPids(pid_t parentPid)
    {
        std::string procPath = "/proc";
        DIR *dirPtr = opendir(procPath.c_str());
        std::vector<pid_t> childPids;
        if (!dirPtr)
        {
            LOG_F(ERROR, "failed in open directory: %s\n", procPath.c_str());
            return childPids;
        }
        dirent *childDir;
        while ((childDir = readdir(dirPtr)) != nullptr)
        {
            int dname = std::atoi(childDir->d_name);
            if (dname == 0)
            {
                continue;
            }
            // try
            // {
            // dname = std::stoi(childDir->d_name);
            pid_t pid;
            std::string procName;
            getProcessInfo(dname, pid, procName);
            if (parentPid == pid)
            {
                childPids.emplace_back(dname);
            }
            // }
            // catch (const std::invalid_argument &e)
            // {
            //     continue;
            // }
            // catch (...)
            // {
            //     // LOG_F(ERROR, "exception: %s\n", e.what());
            //     LOG_F(ERROR, "failed in cast %s to int\n", childDir->d_name);
            //     continue;
            // }
        }
        return childPids;
    }

    ProcessNode ProcessUtil::buildProcessTree(PID_T rootPid)
    {
        ProcessNode root;
        std::unordered_set<PID_T> visited;
        buildTreeRecursive(rootPid, root, visited);
        return root;
    }

    std::string ProcessUtil::printProcessTree(const ProcessNode &root, bool cmdIsVisable)
    {
        std::ostringstream oss;
        printTreeRecursive(root, oss, {}, true, cmdIsVisable);
        return oss.str();
    }

    std::string ProcessUtil::printProcessTree(pid_t rootPid, bool cmdIsVisable)
    {
        std::string str = "\n";
        return str + printProcessTree(buildProcessTree(rootPid), cmdIsVisable);
    }

    int ProcessUtil::killAllDescendantsProcess(PID_T pid, bool isKillSelf)
    {
        ProcessNode root = buildProcessTree(pid);
        std::function<void(ProcessNode & node)> killAllDescendant = [&](ProcessNode &node)
        {
            for (auto &childNode : node.children)
            {
                killAllDescendant(childNode);
                killProcess(childNode);
            }
        };
        // 不递归杀死子进程和孙子进程，会导致只杀死了sh进程，而没有杀死tshark进程，在子线程中的fgets也会一直阻塞或一直获取到值，不会停止
        killAllDescendant(root); // 一定要调用刚刚写好的lambda函数emmmm，写完没调用害我debug一个下午了
        if (isKillSelf)
        {
            return killProcess(root);
        }
        return 0;
    }

    int ProcessUtil::killProcess(const ProcessNode &node)
    {
        LOG_F(INFO, "send sigterm to process pid:%d cmdline:%s", node.pid, node.cmdline.c_str());
#ifdef _WIN32
        // 打开指定进程
        HANDLE hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        if (hProcess == nullptr)
        {
            std::cout << "Failed to open process with PID " << pid << ", error: " << GetLastError() << std::endl;
            return -1;
        }

        // 终止进程
        if (!TerminateProcess(hProcess, 0))
        {
            std::cout << "Failed to terminate process with PID " << pid << ", error: " << GetLastError() << std::endl;
            CloseHandle(hProcess);
            return -1;
        }

        // 成功终止进程
        CloseHandle(hProcess);
        return 0;
#else
        // 这里可以加个waitpid回收子进程
        int res = kill(node.pid, SIGTERM);
        waitPid(node.pid);
        return res;
#endif
    }

    bool ProcessUtil::waitPid(pid_t pid)
    {
        int status;
        std::string msg;
        pid_t tmpPid;
        // 返回-1代表出错,可以perror("waitpid")打印错误
        while ((tmpPid = waitpid(pid, &status, WNOHANG | WUNTRACED)) >= 0)
        {
            if (tmpPid == 0)
            {
                continue;
            }
            if (WIFEXITED(status)) // 子进程正常退出
            {
                msg = fmt::format("exit with status: {}", WEXITSTATUS(status));
            }
            if (WIFSIGNALED(status)) // 子进程由信号终止
            {
                msg = fmt::format("terminate by signum: {}", WTERMSIG(status));
            }
            if (WIFSTOPPED(status)) // 子进程由信号停止
            {
                msg = fmt::format("stop by signum: {}", WSTOPSIG(status));
            }
            LOG_F(WARNING, "child process pid %d: %s\n", pid, msg.c_str());
        }
        return true;
    }

    void ProcessUtil::buildTreeRecursive(PID_T pid, ProcessNode &node, std::unordered_set<PID_T> &visited)
    {
        if (visited.find(pid) != visited.end())
        {
            return;
        }
        visited.insert(pid);
        node.pid = pid;
        getProcessInfo(pid, node.parent_pid, node.name);
        getProcessCmdline(pid, node.cmdline);
        auto childPids = getChildPids(pid);
        for (pid_t childPid : childPids)
        {
            ProcessNode childNode;
            buildTreeRecursive(childPid, childNode, visited);
            // childNode.children.push_back(childNode);
            node.children.push_back(childNode);
        }
    }

    void ProcessUtil::printTreeRecursive(const ProcessNode &node, std::ostream &os, const std::vector<bool> &lastFlags, bool isLast, bool cmdIsVisable)
    {
        // note: 打印进程树
        // 生成前缀
        std::string prefix;
        for (size_t i = 0; i < lastFlags.size(); ++i)
        {
            // 除根节点外，所有层级都根据父级状态显示竖线
            prefix += lastFlags[i] ? "    " : "│   ";
        }

        // 添加当前节点连接符（根节点特殊处理）
        if (!lastFlags.empty())
        {
            prefix += isLast ? "└── " : "├── ";
        }
        else
        {
            prefix = "│── "; // 强制根节点前显示竖线
        }

        // 打印当前节点
        os << prefix
           << node.name
           << " (PID:" << node.pid
           << ") [Parent:" << node.parent_pid << "]";

        if (cmdIsVisable && !node.cmdline.empty())
        {
            os << " ->Cmd: " << node.cmdline << "\n";
        }
        else
        {
            os << "\n";
        }

        // 递归子节点
        for (size_t i = 0; i < node.children.size(); ++i)
        {
            bool childIsLast = (i == node.children.size() - 1);
            std::vector<bool> newFlags = lastFlags;
            newFlags.push_back(isLast);
            printTreeRecursive(node.children[i], os, newFlags, childIsLast, cmdIsVisable);
        }
    }
}